# 实现滑移相关的特征提取，包括特征应变率，粘结滑移率
import numpy as np
import time


def get_C_strains(displacement: np.ndarray, normals: np.ndarray) -> np.ndarray:
    """
    计算特征应变
    Args:
        displacement (np.ndarray): 形状为 (20, 20, 3) 或 (N, 20, 20, 3) 的位移场
        normals (np.ndarray): 形状为 (20, 20, 3) 或 (N, 20, 20, 3) 的法向量场
    """
    if (
        displacement.shape[1:] != (400, 3)
        and displacement.shape[1:] != (20, 20, 3)
        and displacement.shape != (20, 20, 3)
        and displacement.shape != (400, 3)
    ):
        raise ValueError(
            "Position must be of shape (N, 400, 3) or (N, 20, 20, 3) or (400,3) or (20,20,3)"
        )
    if (
        normals.shape[1:] != (400, 3)
        and normals.shape[1:] != (20, 20, 3)
        and normals.shape != (20, 20, 3)
        and normals.shape != (400, 3)
    ):
        raise ValueError(
            "Normals must be of shape (N, 400, 3) or (N, 20, 20, 3) or (400,3) or (20,20,3)"
        )
    
    if displacement.shape == (400, 3):
        displacement = displacement.reshape((20, 20, 3))
    if displacement.shape[1:] == (400, 3):
        displacement = displacement.reshape((-1, 20, 20, 3))
    if normals.shape == (400, 3):
        normals = normals.reshape((20, 20, 3))
    if normals.shape[1:] == (400, 3):
        normals = normals.reshape((-1, 20, 20, 3))

    u_x = displacement[..., 0]
    u_y = displacement[..., 1]
    assert u_x.shape[-2:] == (20, 20)
    u_xx, u_xy = np.gradient(u_x, axis=(-2, -1))
    u_yx, u_yy = np.gradient(u_y, axis=(-2, -1))

    e_xx = u_xx + (u_xx**2 + u_yx**2) / 2
    e_yy = u_yy + (u_xy**2 + u_yy**2) / 2
    e_xy = (u_xy + u_yx) / 2 + (u_xx * u_xy + u_yx * u_yy) / 2

    # 将坐标轴应变旋转为主轴应变
    # i.e. for matrix [a b; b c], calculate its eigenvalues:
    # (a-x)(c-x)-b^2=0
    # ac - (a+c)x + x^2 - b^2=0
    # x^2 - (a+c)x + (ac-b^2)=0
    # x1+x2 = a+c, x1*x2 = ac-b^2
    # x1**2+x2**2 = (x1+x2)**2 - 2*x1*x2 = (a+c)**2 - 2(ac-b^2) = a^2+c^2+2b^2
    # Delta = np.sqrt((e_xx - e_yy) ** 2 + 4 * e_xy**2) / 2
    # e1 = (e_xx + e_yy) / 2 + Delta
    # e2 = (e_xx + e_yy) / 2 - Delta

    # # 计算特征应变，特征应变率后续再进行计算
    # C_strain = np.sqrt(e1**2 + e2**2)
    C_strain = np.sqrt(e_xx**2 + e_yy**2 + 2 * e_xy**2)
    normal_compensate_factor = 1 / normals[..., 2]
    C_strain *= normal_compensate_factor
    return C_strain.reshape(-1,400).squeeze()
